# -*- coding: utf-8 -*-
from typing import List
from textwrap import dedent
from airflow import DAG
from utils.hooks.FeiShuHook import FeiShuHook
from airflow.models import Variable
import re

# 红色：Red => 失败 Red
# 绿色：Green => 成功 Green
# 橙色：Orange => 上游失败 Orange
# 黄色：Yellow => 异常，准备重试 Yellow
# 粉色：default => 跳过 default
dag_template = """\
{% if dag_run.state.upper() == 'FAILED' %}【<font color=Red>{{ dag_run.state.upper() }}</font>】\
{% elif dag_run.state.upper() == 'SUCCESS' %}【<font color=Green>{{ dag_run.state.upper() }}</font>】\
{% else %}【{{ dag_run.state.upper() }}】{% endif %}{{ dag_run.dag_id }} ({{ var.value.env.upper() }} 环境)
**任务状态:** {{ dag_run.state.upper() }}
**任务日期:** {{ execution_date | cst_ds }}
**运行时长:** {{ dag_run.end_date - dag_run.start_date }}
**开始时间:** {{ dag_run.start_date | to_cst }}
**结束时间:** {{ dag_run.end_date | to_cst }}
**成功任务:** {{ dag_run.get_task_instances(state=['success']) | length }} 
**跳过任务:** {{ dag_run.get_task_instances(state=['skipped']) | length }} 
**失败任务:** {{ dag_run.get_task_instances(state=['failed', 'upstream_failed']) | length }} 
{% for ti in dag_run.get_task_instances(state=['failed']) %}\
  > <font color=red>{{ ti.task_id }}</font> 
{% endfor %}
{% for ti in dag_run.get_task_instances(state=['upstream_failed']) %}\
  > <font color=orange>{{ ti.task_id }}(上游失败)</font> 
{% endfor %}
"""


def feishu_dag_alert(feishu_conn_id: str = 'feishu_default', at_mobiles: List[str] = None, at_all: bool = False):
    def _feishu_dag_alert(context, *args, **kwargs):
        jinja_env = context['dag'].get_template_env()
        runStatus = jinja_env.from_string("{{ dag_run.state.upper() }}").render(context)
        if runStatus == 'FAILED':
            level = 'red'
        else:
            level = 'green'
        title = jinja_env.from_string("【{{ dag_run.state.upper() }}】{{ dag_run.dag_id }} ({{ var.value.env.upper() }} 环境)").render(context)
        print("context:",str(context))
        content = dedent(jinja_env.from_string(dag_template).render(context))
        print("content:",content)
        return FeiShuHook(feishu_conn_id=feishu_conn_id,
                          message_type='markdown',
                          title=title,
                          message=content,
                          at_mobiles=at_mobiles,
                          at_all=at_all,
                          level=level).send()

    return _feishu_dag_alert


task_template = """
任务 {{ ti.task_id }} {% if ti.state.upper() == 'FAILED' %}<font color=Red>{{ ti.state.upper() }}</font>\
{% elif ti.state.upper() == 'SUCCESS' %}<font color=Green>{{ ti.state.upper() }}</font>\
{% elif ti.state.upper() == 'UP_FOR_RETRY' %}<font color=Yellow>{{ ti.state.upper() }}</font>\
{% elif ti.state.upper() == 'SKIPPED' %}<font color=default>{{ ti.state.upper() }}</font>\
{% else %}{{ ti.state.upper() }}{% endif %} ({{ var.value.env.upper() }} 环境)

**任务日期:** {{ execution_date | cst_ds }}
**开始时间:** {{ ti.start_date | to_cst }}
{% if exception | trim | length %}\
**异常信息:**
  > {{ exception }}
{% endif %}
"""


def feishu_task_alert(feishu_conn_id: str = 'feishu_default', at_mobiles: List[str] = None, at_all: bool = False):
    def _feishu_task_alert(context, *args, **kwargs):
        jinja_env = context['dag'].get_template_env()
        runStatus = jinja_env.from_string("{{ ti.state.upper() }}").render(context)
        if runStatus == 'SUCCESS':
            level = 'green'
        elif runStatus=='SKIPPED':
            level = 'grey'
        elif runStatus=='FAILED':
            level='red'
        else:
            level='yellow'
        title = jinja_env.from_string("任务 {{ ti.task_id }} {{ ti.state.upper() }} ({{ var.value.env.upper() }} 环境)").render(context)
        content = dedent(jinja_env.from_string(task_template).render(context))
        contentnew = re.sub(r'"password"((?!"password").)*"driver"', '"password":"******",\n"driver"', content, flags=re.DOTALL)
        print(f'send to feishu: {feishu_conn_id}')
        print(f'title: \n{title}')
        print(f'content: \n{contentnew}')

        return FeiShuHook(feishu_conn_id=feishu_conn_id,
                          message_type='markdown',
                          title=title,
                          message=contentnew,
                          at_mobiles=at_mobiles,
                          at_all=at_all,
                          level=level).send()

    return _feishu_task_alert


def feishu_sla_alert(feishu_conn_id: str = 'feishu_default', at_mobiles: List[str] = None, at_all: bool = False):
    def _feishu_sla_alert(dag: DAG, task_list: str, blocking_task_list: str, slas, blocking_tis):
        env = Variable.get('env', '').upper()
        title = f'告警：任务有超时风险[{env}环境]'
        content = f'以下任务有<font color=Yellow>超时</font>风险 ({ env } 环境)\n\n' + '\n'.join(f'- {sla.task_id}' for sla in slas)

        print(f'send to feishu: {feishu_conn_id}')
        print(f'title: \n{title}')
        print(f'content: \n{content}')

        return FeiShuHook(feishu_conn_id=feishu_conn_id,
                          message_type='markdown',
                          title=title,
                          message= content,
                          at_mobiles=at_mobiles,
                          at_all=at_all,
                          level='yellow' ).send()

    return _feishu_sla_alert
